from model import BertSimMatchModel
from flask import Flask, request
from flask_cors import CORS

import time
import psutil
import jieba
import pynvml

pynvml.nvmlInit()

MODEL_DIR = 'model'
model = BertSimMatchModel.load(MODEL_DIR)  # 加载 Bert 相似匹配模型

app = Flask(__name__)
CORS(app)


@app.route('/judge', methods=["GET"])
def judge():
    question = request.args.get("question")
    right_answer = request.args.get("rightAnswer")
    user_answer = request.args.get("userAnswer")

    start_time = time.time()
    results = model.predict([(question, right_answer, user_answer)])[0]
    end_time = time.time()

    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)

    return {
        "label": results[0],
        "prob": results[1],
        "time": end_time-start_time,
        "RAM": psutil.virtual_memory().percent,
        "cpu": psutil.cpu_percent(None),
        "gpu": meminfo.used,
        "question": jieba.lcut(question),
        "rightAnswer": jieba.lcut(right_answer),
        "userAnswer": jieba.lcut(user_answer)
    }

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0')
